/*
 * PyEigenMatrix.cpp
 *
 *  Created on: 6 Feb 2012
 *      Author: allan
 */

#include "PyEigenMatrix.h"

// Eigen includes
#include <Eigen/Core>
using namespace Eigen;

// Boost includes
#include <boost/python.hpp>
#include <boost/range.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/python/slice.hpp>
using namespace boost::python;

// PyGeoReact includes
#define GEOREACT_NUMPY_INTERNAL
#include <Common/NumpyArray.h>
#include <Common/PyErrors.h>

template<typename Scalar, int NumpyType>
struct WrapperEigenMatrix
{
	typedef Matrix<Scalar, -1, -1> MatrixType;

	static boost::shared_ptr<MatrixType>
	init_default()
	{
		return boost::shared_ptr<MatrixType>(new MatrixType());
	}
	
	static boost::shared_ptr<MatrixType>
	init_copy(const MatrixType& copy)
	{
		return boost::shared_ptr<MatrixType>(new MatrixType(copy));
	}
	
	static boost::shared_ptr<MatrixType>
	init_with_rows_cols(int rows, int cols)
	{
		return boost::shared_ptr<MatrixType>(new MatrixType(rows, cols));
	}
	
	static boost::shared_ptr<MatrixType>
	init_with_rows_cols_and_value(int rows, int cols, const Scalar& val)
	{
		boost::shared_ptr<MatrixType> mat = init_with_rows_cols(rows, cols);
		
		std::fill_n(mat->data(), rows * cols, val);
		
		return mat;
	}
	
	static boost::shared_ptr<MatrixType>
	init_with_array(const numeric::array& array)
	{
		tuple shape = extract<tuple>(array.attr("shape"));
		
		long rows, cols;
		
		if(len(shape) == 1)
		{
			rows = 1;
			cols = extract<long>(shape[0]);
		}
		else
		{
			rows = extract<long>(shape[0]);
			cols = extract<long>(shape[1]);
		}

		boost::shared_ptr<MatrixType> mat = init_with_rows_cols(rows, cols);
		
		if(PyArray_TYPE(array.ptr()) != NumpyType)
			TypeError("type mismatch between array and matrix entries");
		
		Scalar* data = (Scalar*) PyArray_DATA(array.ptr());
		
		int k = 0;
		for(int i = 0; i < rows; ++i)
			for(int j = 0; j < cols; ++j)
				mat->operator()(i, j) = data[k++];
		
		// --- does not work because Eigen is col-major ---
		// std::copy(data, data + rows * cols, mat->data());
		
		return mat;
	}
	
	template<typename Sequence>
	static boost::shared_ptr<MatrixType>
	init_with_sequence(const Sequence& seq)
	{
		return init_with_array(numeric::array(seq));
	}
	
	static void
	get_slice_data(MatrixType& self, const object& arg, long& start, long& stop, long max)
	{
		extract<slice> xslice(arg);
		
		if(!xslice.check())
		{
			start = extract<long>(arg);
			stop = start + 1;
		}
		else
		{
			slice slice = xslice();
			if(!slice.step().is_none())
				SliceStepError();
			
			start = 0;
			stop = max;
			
			if(!slice.start().is_none())
				start = extract<long>(slice.start());
			
			if(!slice.stop().is_none())
				stop = extract<long>(slice.stop());
		}
		
		if(start < 0 || stop < start || stop > max)
			IndexError();
	}
	
	static void
	set_item(MatrixType& self, const tuple& tuple, const Scalar& scalar)
	{
		long row_start, row_stop;
		long col_start, col_stop;
		
		get_slice_data(self, tuple[0], row_start, row_stop, self.rows());
		get_slice_data(self, tuple[1], col_start, col_stop, self.cols());
		
		if(row_stop - row_start == 1 && col_stop - row_start == 1)
			self(row_start, col_start) = scalar;
		else
			for(int i = row_start; i < row_stop; ++i)
				for(int j = col_start; j < col_stop; ++j)
					self(i, j) = scalar;
	}
	
	static object
	get_item(MatrixType& self, const tuple& tuple)
	{
		long row_start, row_stop;
		long col_start, col_stop;
		
		get_slice_data(self, tuple[0], row_start, row_stop, self.rows());
		get_slice_data(self, tuple[1], col_start, col_stop, self.cols());
		
		const long row_size = row_stop - row_start;
		const long col_size = col_stop - col_start;
		
		if(row_size == 1 && col_size == 1)
			return object(self(row_start, col_start));
		else
			return object(MatrixType(self.block(row_start, col_start, row_size, col_size)));
	}
	
	static void
	set_slice_with_eigen_matrix(MatrixType& self, const tuple& tuple, const MatrixType& mat)
	{
		long row_start, row_stop;
		long col_start, col_stop;
		
		get_slice_data(self, tuple[0], row_start, row_stop, self.rows());
		get_slice_data(self, tuple[1], col_start, col_stop, self.cols());
		
		if(mat.rows() != row_stop - row_start || mat.cols() != col_stop - col_start)
			IndexError();
		
		self.block(row_start, col_start, mat.rows(), mat.cols()) = mat;
	}
	
	static void
	set_slice_with_array(MatrixType& self, const tuple& tuple, const numeric::array& array)
	{
		long row_start, row_stop;
		long col_start, col_stop;
		
		get_slice_data(self, tuple[0], row_start, row_stop, self.rows());
		get_slice_data(self, tuple[1], col_start, col_stop, self.cols());
		
		long rows = row_stop - row_start;
		long cols = col_stop - col_start;
		
		if(PyArray_TYPE(array.ptr()) != NumpyType)
			TypeError("type mismatch between array and matrix entries");
		
		Scalar* data = (Scalar*) PyArray_DATA(array.ptr());
		
		int k = 0;
		for(int i = row_start; i < rows; ++i)
			for(int j = col_start; j < cols; ++j)
				self(i, j) = data[k++];
		
		// --- does not work because Eigen is col-major ---
		// auto self_data = self.block(row_start, col_start, rows, cols).data();
		// std::copy(data, data + rows * cols, self_data);
	}
	
	template<typename Sequence>
	static void
	set_slice_with_sequence(MatrixType& self, const tuple& tuple, const Sequence& seq)
	{
		set_slice_with_array(self, tuple, numeric::array(seq));
	}
	
	static numeric::array
	array(MatrixType& self)
	{
		const int rows = self.rows();
		const int cols = self.cols();
		
		npy_intp dims[2] = {rows, cols};
		
		// --- does not work because Eigen is col-major ---
		// object obj(handle<>(
		//	incref(PyArray_SimpleNewFromData(2, dims, NumpyType, map_self.data()))));
		
		object obj(handle<>(incref(PyArray_SimpleNew(2, dims, NumpyType))));
		
		Scalar* data = (Scalar*) PyArray_DATA(obj.ptr());
		
		int k = 0;
		for(int i = 0; i < rows; ++i)
			for(int j = 0; j < cols; ++j)
				data[k++] = self(i, j);
		
		return extract<numeric::array>(obj);
	}
	
	static Scalar*
	begin(MatrixType& self)
	{
		return self.data();
	}
	
	static Scalar*
	end(MatrixType& self)
	{
		return self.data() + self.rows() * self.cols();
	}
};

template<typename Scalar, int NumpyType>
void
export_EigenMatrixType(const char* class_name)
{
	typedef Matrix<Scalar, -1, -1> MatrixType;
	typedef WrapperEigenMatrix<Scalar, NumpyType> WrapperType;
	
	class_<MatrixType, boost::shared_ptr<MatrixType>>(class_name).def(
		"__init__", make_constructor(&WrapperType::init_default))

	.def("__init__", make_constructor(&WrapperType::init_copy))

	.def("__init__", make_constructor(&WrapperType::init_with_rows_cols))

	.def("__init__", make_constructor(&WrapperType::init_with_rows_cols_and_value))

	.def("__init__", make_constructor(&WrapperType::init_with_array))

	.def("__init__", make_constructor(&WrapperType::template init_with_sequence<list>))

	.def("__init__", make_constructor(&WrapperType::template init_with_sequence<tuple>))

	.def("__len__", &MatrixXd::rows)

	.def("__setitem__", &WrapperType::set_item, with_custodian_and_ward<1, 2>()) // to let container keep value
	
	.def("__getitem__", &WrapperType::get_item)

	.def("__setitem__", &WrapperType::set_slice_with_eigen_matrix)

	.def("__setitem__", &WrapperType::set_slice_with_array)

	.def("__setitem__", &WrapperType::template set_slice_with_sequence<list>)

	.def("__setitem__", &WrapperType::template set_slice_with_sequence<tuple>)

	.def("__iter__", range(&WrapperType::begin, &WrapperType::end))

	.def("__array__", &WrapperType::array)

	.add_property("rows", &MatrixXd::rows)

	.add_property("cols", &MatrixXd::cols)

	.def("array", &WrapperType::array)

	.def(self_ns::str(self_ns::self));
}

void
export_EigenMatrix()
{
	export_EigenMatrixType<double, NPY_DOUBLE>("MatrixXd");
	export_EigenMatrixType<float, NPY_FLOAT>("MatrixXf");
	export_EigenMatrixType<int, NPY_INT>("MatrixXi");
}
